import copy
from ...utils import BaseModuleBuilder

from .fcn import FCN
from .fcnbigmodel import FCNBigModel
from .mask2former import Mask2Former
from .maskformer import MaskFormer
from .sct import SCT
from .grnet import GRNet
from .ddrnet import DDRNet
from .deeplabv3plus import Deeplabv3Plus
from .deeplabv3plusbigmodel import Deeplabv3PlusBigModel
'''SegmentorBuilder'''
class SegmentorBuilder(BaseModuleBuilder):
    REGISTERED_MODULES = {
        'FCN': FCN, 'Mask2Former': Mask2Former, 'MaskFormer': MaskFormer,'SCT':SCT,
        'GRNet':GRNet, 'DDRNet':DDRNet,'Deeplabv3Plus':Deeplabv3Plus, 'FCNBigModel':FCNBigModel,'Deeplabv3PlusBigModel':Deeplabv3PlusBigModel
    }
    '''build'''
    def build(self, segmentor_cfg, mode):
        segmentor_cfg = copy.deepcopy(segmentor_cfg)
        segmentor_type = segmentor_cfg.pop('type')
        segmentor = self.REGISTERED_MODULES[segmentor_type](cfg=segmentor_cfg, mode=mode)
        return segmentor


'''BuildSegmentor'''
BuildSegmentor = SegmentorBuilder().build
